from ..models import *
from typing import List, Optional, Dict

__all__ = [
    "OpenaiChatCompletionRequest",
    "OpenaiChatCompletionResponse",
    "OpenaiChoice",
    "OpenaiCompletionUsage",
]


class OpenaiChatCompletionRequest(BaseModel):
    messages: List[Dict] = Field(
        ...,
        description="A list of messages comprising the conversation so far. \
        [Example Python code](https://cookbook.openai.com/examples/how_to_format_inputs_to_chatgpt_models).",
    )

    model: str = Field(..., description="`model` is a valid TaskingAI model ID or assistant ID.")

    function_call: Optional[OpenaiFunctionCallOption] = Field(
        None,
        description="Deprecated in favor of `tool_choice`. Controls which (if any) function is called by the model.\
        `none` means the model will not call a function and instead generates a message. `auto` means the model\
        can pick between generating a message or calling a function. Specifying a particular function via\
        `{'name': 'my_function'}` forces the model to call that function. `none` is the default when no functions\
        are present. `auto` is the default if functions are present.",
    )

    functions: Optional[List[OpenaiFunction]] = Field(
        None, description="Deprecated in favor of `tools`. A list of functions the model may generate JSON inputs for."
    )

    tool_choice: Optional[OpenaiFunctionCallOption] = Field(
        None,
        description="Controls which (if any) function is called by the model. `none` means the model will not call\
        a function and instead generates a message. `auto` means the model can pick between generating a message or\
        calling a function. Specifying a particular function via `{'type': 'function', 'function': {'name':\
        'my_function'}}` forces the model to call that function. `none` is the default when no functions are\
        present. `auto` is the default if functions are present.",
    )

    tools: Optional[List[OpenaiTool]] = Field(
        None,
        description="A list of tools the model may call. Currently, only functions are supported as a tool.\
        Use this to provide a list of functions the model may generate JSON inputs for. A max of 128 functions are\
        supported.",
    )

    stream: bool = Field(
        False,
        description="Indicates whether the response should be streamed. If set to True, the response will be streamed\
        using Server-Sent Events (SSE).",
    )


class OpenaiChoice(BaseModel):
    finish_reason: Literal["stop", "length", "tool_calls", "content_filter", "function_call"] = Field(
        ...,
        description="The reason the model stopped generating tokens. This will be stop if the model hit a natural stop point or a provided stop sequence, length if the maximum number of tokens specified in the request was reached, content_filter if content was omitted due to a flag from our content filters, tool_calls if the model called a tool, or function_call (deprecated) if the model called a function.",
    )

    index: int = Field(..., description="The index of the choice in the list of choices.")

    logprobs: None

    message: OpenaiChatCompletionMessage = Field(..., description="A chat completion message generated by the model.")


class OpenaiCompletionUsage(BaseModel):
    completion_tokens: int = Field(..., description="Number of tokens in the generated completion.")
    prompt_tokens: int = Field(..., description="Number of tokens in the prompt.")
    total_tokens: int = Field(..., description="Total number of tokens used in the request (prompt + completion).")


class OpenaiChatCompletionResponse(BaseModel):
    id: str = Field(..., description="A unique identifier for the chat completion.")
    choices: List[OpenaiChoice] = Field(
        ..., description="A list of chat completion choices. Can be more than one if `n` is greater than 1."
    )
    created: int = Field(..., description="The Unix timestamp (in seconds) of when the chat completion was created.")
    model: str = Field(..., description="The model used for the chat completion.")
    object: Literal["chat.completion"] = Field(..., description="The object type, which is always `chat.completion`.")
    system_fingerprint: Optional[str] = Field(
        default=None,
        description="This fingerprint represents the backend configuration that the model runs with. Can be used in conjunction with the `seed` request parameter to understand when backend changes have been made that might impact determinism.",
    )
    usage: Optional[OpenaiCompletionUsage] = Field(
        default=None, description="Usage statistics for the completion request."
    )
